#include "postprocess/kalman_filter/motion_kalmanfilter.h"

MotionKalmanFilter::MotionKalmanFilter() {
  // x y vx vy ax ay
}

MotionKalmanFilter::MotionKalmanFilter(const ProcessorClass& flag) {
  processor_flag_ = flag;
}

MotionKalmanFilter::~MotionKalmanFilter() {}

void MotionKalmanFilter::Init(int dynam_param_num, int measure_param_num,
                              int control_param_num) {
  BaseKalmanFilter::Init(dynam_param_num, measure_param_num, control_param_num);
  age_ = 0;
  P_.block<2, 2>(0, 0) *= initial_vehicle_position_noise;
  P_.block<2, 2>(2, 2) *= initial_vehicle_velocity_noise;
  P_.block<2, 2>(4, 4) *= initial_vehicle_acceleration_noise;
  H_.setIdentity();

  // 初始化Q阵 状态量误差
  Q_.setIdentity(dynam_param_num, dynam_param_num);
  // Q_ *= initial_env_noiose;

  // 初始化R阵
  R_.setIdentity(measure_param_num, measure_param_num);
}

void MotionKalmanFilter::SetCenterPoint(const Eigen::Vector3f& point) noexcept {
  belief_center_point_ = point;
  State_.head<2>() = point.head<2>();
}

void MotionKalmanFilter::SetVelocity(const Eigen::Vector3f& vel) noexcept {
  belief_velocity_ = vel;
  State_(2) = vel.x();
  State_(3) = vel.y();
}

void MotionKalmanFilter::SetAcceleration(const Eigen::Vector3f& acc) noexcept {
  belief_acceleration_ = acc;
  State_.tail<2>() = acc.head<2>();
}

void MotionKalmanFilter::SetSettings(
    PostProcessorSettings const& settings) noexcept {
  settings_ = settings;
}

void MotionKalmanFilter::GetCenterPoint(Eigen::Vector3f& point) noexcept {
  point = belief_center_point_;
}

void MotionKalmanFilter::GetVelocity(Eigen::Vector3f& vel) noexcept {
  vel = belief_velocity_;
}

void MotionKalmanFilter::GetAcceleration(Eigen::Vector3f& acc) noexcept {
  acc = belief_acceleration_;
}

void MotionKalmanFilter::GetProcessorClassFlag(ProcessorClass& flag) noexcept {
  flag = processor_flag_;
}

void MotionKalmanFilter::SetProcessorClassFlag(ProcessorClass flag) noexcept {
  processor_flag_ = flag;
}

void MotionKalmanFilter::SetQ(Eigen::MatrixXf& Q) { Q_ = Q; }

void MotionKalmanFilter::SetR(Eigen::MatrixXf& R) { R_ = R; }

void MotionKalmanFilter::GetState(Eigen::Vector3f* center_point, float& vx,
                                  float& vy, float& ax, float& ay) noexcept {
  *center_point = belief_center_point_;
  vx = belief_velocity_.x();
  vy = belief_velocity_.y();
  ax = belief_acceleration_.x();
  ay = belief_acceleration_.y();
}

void MotionKalmanFilter::GetState(Eigen::Vector3f* center_point,
                                  Eigen::Vector3f* velocity,
                                  Eigen::Vector3f* acceleration) noexcept {
  (*center_point) = belief_center_point_;
  (*velocity) = belief_velocity_;
  (*acceleration) = belief_acceleration_;
}

void MotionKalmanFilter::UpdateWithObject(TrackedObjectPtr& new_object,
                                          TrackedObjectPtr const& old_object,
                                          double time_diff) {
  // kalman filter update
  update_state(new_object, old_object, time_diff);
  age_ += 1;
}

void MotionKalmanFilter::UpdateWithoutObject(double time_diff) {
  belief_center_point_ += belief_velocity_ * time_diff;
  age_ += 1;
}

void MotionKalmanFilter::update_state(TrackedObjectPtr const& new_object,
                                      TrackedObjectPtr const& old_object,
                                      double time_diff) {
  // 利用new_object中的processor_flag来进行判别
  //  A赋值
  A_(0, 2) = time_diff;
  A_(1, 3) = time_diff;
  A_(2, 4) = time_diff;
  A_(3, 5) = time_diff;
  A_(0, 4) = time_diff * time_diff * 0.5;
  A_(1, 5) = time_diff * time_diff * 0.5;

  if (processor_flag_ == ProcessorClass::Vehicle) {
    Q_(0, 0) = settings_.motion_vehicle_Q[0];
    Q_(1, 1) = settings_.motion_vehicle_Q[7];
    Q_(2, 2) = settings_.motion_vehicle_Q[14];
    Q_(3, 3) = settings_.motion_vehicle_Q[21];
    Q_(4, 4) = settings_.motion_vehicle_Q[28];
    Q_(5, 5) = settings_.motion_vehicle_Q[35];

    R_(0, 0) = settings_.motion_vehicle_R[0];
    R_(1, 1) = settings_.motion_vehicle_R[5];
    R_(2, 2) = settings_.motion_vehicle_R[10];
    R_(3, 3) = settings_.motion_vehicle_R[15];

  } else if (processor_flag_ == ProcessorClass::Pedestrian) {
    Q_(0, 0) = settings_.motion_pedestrian_Q[0];
    Q_(1, 1) = settings_.motion_pedestrian_Q[7];
    Q_(2, 2) = settings_.motion_pedestrian_Q[14];
    Q_(3, 3) = settings_.motion_pedestrian_Q[21];
    Q_(4, 4) = settings_.motion_pedestrian_Q[28];
    Q_(5, 5) = settings_.motion_pedestrian_Q[35];

    R_(0, 0) = settings_.motion_pedestrian_R[0];
    R_(1, 1) = settings_.motion_pedestrian_R[5];
    R_(2, 2) = settings_.motion_pedestrian_R[10];
    R_(3, 3) = settings_.motion_pedestrian_R[15];
  }

  // to do
  // Assign measurement value
  Eigen::VectorXf Z = Eigen::VectorXf::Zero(measure_param_num_);
  Z(0) = new_object->center.x();
  Z(1) = new_object->center.y();
  Z(2) = new_object->velocity.x();
  Z(3) = new_object->velocity.y();

  // std::cout << "--------------------kalmanfilter-----------------------"
  //           << std::endl;
  // std::cout << "previous data : " << std::endl;
  // std::cout << "object center : " << new_object->center.x() << " "
  //           << new_object->center.y() << " vel = " <<
  //           new_object->velocity.x()
  //           << " " << new_object->velocity.y()
  //           << " size = " << new_object->size.x() << " " <<
  //           new_object->size.y()
  //           << std::endl;
  // std::cout << "state center : " << belief_center_point_[0] << " "
  //           << belief_center_point_[1] << " vel = " << belief_velocity_[0]
  //           << " " << belief_velocity_[1] << std::endl;

  // update
  // x = A * x + B * u
  State_ = A_ * State_;
  // p = A * p * A^T + Q
  P_ = A_ * P_ * A_.transpose() + Q_;
  // K = p * H^T / (H * p * H^T + R)
  auto K = P_ * H_.transpose() * (H_ * P_ * H_.transpose() + R_).inverse();
  // x= x + K * (z - H * x)
  Eigen::VectorXf state_gain;
  state_gain.resize(6);
  state_gain =
      static_cast<Eigen::Matrix<float, 6, 1, 0, 6, 1>>(K * (Z - H_ * State_));
  // state_gain_adjustment(state_gain);
  State_ = State_ + state_gain;
  // p = (I - K * H) * p
  Eigen::MatrixXf I =
      Eigen::MatrixXf::Identity(dynam_param_num_, dynam_param_num_);
  P_ = (I - K * H_) * P_;

  // assign value
  belief_center_point_.head<2>() = State_.head<2>();
  belief_center_point_[2] = new_object->center.z();  // z值不更新，直接赋值
  belief_velocity_.head<2>() = State_.segment(2, 2);
  belief_velocity_(2) = 0.0f;
  belief_acceleration_.head<2>() = State_.segment(4, 2);
  belief_acceleration_(2) = 0.0f;

  // std::cout << "after flitering data : " << std::endl;
  // std::cout << "center : " << belief_center_point_[0] << " "
  //           << belief_center_point_[1] << " vel = " << belief_velocity_[0]
  //           << " " << belief_velocity_[1] << std::endl;

  // std::cout << "---------------------------------------------------------"
  //           << std::endl;
}

void MotionKalmanFilter::state_gain_adjustment(Eigen::VectorXf& gain,
                                               double time_diff) {
  if (settings_.use_kalman_breakdown) {
    gain.resize(6);
    if (processor_flag_ == ProcessorClass::Vehicle) {
      // 距离限幅
      float distance_gain = gain.head<2>().norm();
      // float vehicle_distance_breakdown_threshold =
      //     std::min(settings_.vehicle_distance_breakdown_threshold,
      //              static_cast<float>(time_diff) *
      //                  settings_.vehicle_velocity_breakdown_threshold);
      if (distance_gain > settings_.vehicle_distance_breakdown_threshold) {
        gain.head<2>() *=
            settings_.vehicle_distance_breakdown_threshold / distance_gain;
      }

      // 速度限幅
      float velocity_gain = gain.segment(2, 2).norm();
      if (velocity_gain > settings_.vehicle_velocity_breakdown_threshold) {
        gain.segment(2, 2) *=
            settings_.vehicle_velocity_breakdown_threshold / velocity_gain;
      }

      // 加速度限幅
      float accleration_gain = gain.tail<2>().norm();
      if (accleration_gain >
          settings_.vehicle_acceleration_breakdown_threshold) {
        gain.tail<2>() *= settings_.vehicle_acceleration_breakdown_threshold /
                          accleration_gain;
      }
    } else if (processor_flag_ == ProcessorClass::Pedestrian) {
      gain.resize(6);
      // 距离限幅
      float distance_gain = gain.head<2>().norm();
      // float pedestrian_distance_breakdown_threshold =
      //     std::min(settings_.pedestrian_distance_breakdown_threshold,
      //              static_cast<float>(time_diff) *
      //                  settings_.pedestrian_velocity_breakdown_threshold);
      if (distance_gain > settings_.pedestrian_distance_breakdown_threshold) {
        gain.head<2>() *=
            settings_.pedestrian_distance_breakdown_threshold / distance_gain;
      }

      // 速度限幅
      float velocity_gain = gain.segment(2, 2).norm();
      if (velocity_gain > settings_.pedestrian_velocity_breakdown_threshold) {
        gain.segment(2, 2) *=
            settings_.pedestrian_velocity_breakdown_threshold / velocity_gain;
      }

      // 加速度限幅
      float accleration_gain = gain.tail<2>().norm();
      if (accleration_gain >
          settings_.pedestrian_acceleration_breakdown_threshold) {
        gain.tail<2>() *=
            settings_.pedestrian_acceleration_breakdown_threshold /
            accleration_gain;
      }
    }
  }
}
